﻿using System;
using System.Collections.Generic;
using System.Text;

namespace MLForgeSharp.Models.SupervisedLearningModels.EnsembleLearning.Bagging
{
    /// <summary>
    /// 随机森林
    /// </summary>
    public class RandomForest
    {
        private class TreeNode
        {
            public int SplitFeature { get; set; } // 分裂特征
            public double SplitValue { get; set; } // 分裂值
            public TreeNode LeftChild { get; set; } // 左子节点
            public TreeNode RightChild { get; set; } // 右子节点
            public string Label { get; set; } // 叶节点的标签
        }

        private List<TreeNode> forest; // 随机森林
        private int numTrees; // 树的数量
        private int maxFeatures; // 每棵树使用的最大特征数
        private Random random;

        public RandomForest(int numTrees = 100, int maxFeatures = 2)
        {
            this.numTrees = numTrees;
            this.maxFeatures = maxFeatures;
            this.forest = new List<TreeNode>();
            this.random = new Random();
        }

        // 计算基尼不纯度
        private double CalculateGiniImpurity(int[] counts)
        {
            double gini = 1.0;
            int total = 0;
            foreach (int count in counts)
            {
                total += count;
            }
            foreach (int count in counts)
            {
                if (total > 0)
                {
                    double probability = (double)count / total;
                    gini -= probability * probability;
                }
            }
            return gini;
        }

        // 获取标签的计数
        private int[] GetLabelCounts(List<double[]> data, int labelIndex)
        {
            Dictionary<double, int> labelCounts = new Dictionary<double, int>();
            foreach (double[] row in data)
            {
                double label = row[labelIndex];
                if (!labelCounts.ContainsKey(label))
                {
                    labelCounts[label] = 0;
                }
                labelCounts[label]++;
            }

            int[] counts = new int[labelCounts.Count];
            int i = 0;
            foreach (int count in labelCounts.Values)
            {
                counts[i++] = count;
            }
            return counts;
        }

        // 构建决策树
        private TreeNode BuildTree(List<double[]> data, List<int> featureIndices, int labelIndex)
        {
            if (data.Count == 0)
            {
                return new TreeNode { Label = "Unknown" };
            }

            int[] labelCounts = GetLabelCounts(data, labelIndex);
            if (labelCounts.Length == 1)
            {
                return new TreeNode { Label = data[0][labelIndex].ToString() };
            }

            if (featureIndices.Count == 0)
            {
                return new TreeNode { Label = GetMajorityLabel(data, labelIndex) };
            }

            int bestFeatureIndex = 0;
            double bestSplitValue = 0.0;
            double bestGini = double.MaxValue;

            foreach (int featureIndex in featureIndices)
            {
                List<double> featureValues = new List<double>();
                foreach (double[] row in data)
                {
                    featureValues.Add(row[featureIndex]);
                }
                featureValues.Sort();

                for (int i = 1; i < featureValues.Count; i++)
                {
                    double splitValue = (featureValues[i - 1] + featureValues[i]) / 2;
                    List<double[]> leftSubset = new List<double[]>();
                    List<double[]> rightSubset = new List<double[]>();

                    foreach (double[] row in data)
                    {
                        if (row[featureIndex] <= splitValue)
                        {
                            leftSubset.Add(row);
                        }
                        else
                        {
                            rightSubset.Add(row);
                        }
                    }

                    double gini = CalculateGiniImpurity(GetLabelCounts(leftSubset, labelIndex)) * leftSubset.Count / data.Count +
                                  CalculateGiniImpurity(GetLabelCounts(rightSubset, labelIndex)) * rightSubset.Count / data.Count;

                    if (gini < bestGini)
                    {
                        bestGini = gini;
                        bestFeatureIndex = featureIndex;
                        bestSplitValue = splitValue;
                    }
                }
            }

            TreeNode node = new TreeNode
            {
                SplitFeature = bestFeatureIndex,
                SplitValue = bestSplitValue
            };

            List<double[]> leftData = new List<double[]>();
            List<double[]> rightData = new List<double[]>();
            foreach (double[] row in data)
            {
                if (row[bestFeatureIndex] <= bestSplitValue)
                {
                    leftData.Add(row);
                }
                else
                {
                    rightData.Add(row);
                }
            }

            List<int> remainingFeatures = new List<int>(featureIndices);
            remainingFeatures.Remove(bestFeatureIndex);

            node.LeftChild = BuildTree(leftData, remainingFeatures, labelIndex);
            node.RightChild = BuildTree(rightData, remainingFeatures, labelIndex);

            return node;
        }

        // 获取多数标签
        private string GetMajorityLabel(List<double[]> data, int labelIndex)
        {
            Dictionary<double, int> labelCounts = new Dictionary<double, int>();
            foreach (double[] row in data)
            {
                double label = row[labelIndex];
                if (!labelCounts.ContainsKey(label))
                {
                    labelCounts[label] = 0;
                }
                labelCounts[label]++;
            }

            double majorityLabel = 0.0;
            int maxCount = 0;
            foreach (var kvp in labelCounts)
            {
                if (kvp.Value > maxCount)
                {
                    maxCount = kvp.Value;
                    majorityLabel = kvp.Key;
                }
            }
            return majorityLabel.ToString();
        }

        // 训练模型
        public void Train(double[][] data, int labelIndex)
        {
            int numFeatures = data[0].Length - 1;
            for (int i = 0; i < numTrees; i++)
            {
                List<double[]> subSample = new List<double[]>();
                for (int j = 0; j < data.Length; j++)
                {
                    subSample.Add(data[random.Next(data.Length)]);
                }

                List<int> featureIndices = new List<int>();
                while (featureIndices.Count < maxFeatures)
                {
                    int index = random.Next(numFeatures);
                    if (!featureIndices.Contains(index))
                    {
                        featureIndices.Add(index);
                    }
                }

                forest.Add(BuildTree(subSample, featureIndices, labelIndex));
            }
        }

        // 预测
        public string Predict(double[] features)
        {
            Dictionary<string, int> labelCounts = new Dictionary<string, int>();
            foreach (TreeNode tree in forest)
            {
                TreeNode node = tree;
                while (node.Label == null)
                {
                    if (features[node.SplitFeature] <= node.SplitValue)
                    {
                        node = node.LeftChild;
                    }
                    else
                    {
                        node = node.RightChild;
                    }
                }
                if (!labelCounts.ContainsKey(node.Label))
                {
                    labelCounts[node.Label] = 0;
                }
                labelCounts[node.Label]++;
            }

            string majorityLabel = "";
            int maxCount = 0;
            foreach (var kvp in labelCounts)
            {
                if (kvp.Value > maxCount)
                {
                    maxCount = kvp.Value;
                    majorityLabel = kvp.Key;
                }
            }
            return majorityLabel;
        }
    }

    // 示例程序
    public class RandomForestExample
    {
        public RandomForestExample()
        {
            // 示例数据
            double[][] data = new double[][]
        {
            new double[] { 1.0, 2.0, 0.0 },
            new double[] { 1.1, 2.1, 0.0 },
            new double[] { 1.2, 2.2, 0.0 },
            new double[] { 1.3, 2.3, 0.0 },
            new double[] { 10.0, 10.0, 1.0 },
            new double[] { 10.1, 10.1, 1.0 }
        };

            int labelIndex = 2; // 标签列索引

            // 创建模型
            RandomForest model = new RandomForest(numTrees: 100, maxFeatures: 2);

            // 训练模型
            model.Train(data, labelIndex);

            // 预测
            double[] testFeatures = { 1.0, 2.0 };
            string predictedLabel = model.Predict(testFeatures);

            System.Console.WriteLine("预测标签: " + predictedLabel);
        }
    }
}
